from typing import List

from redis import Redis

from .nodes.mcp import MCPNodes
from ..graph.base import BaseGraph
from ..graph.nodes import *
from ..llms import LLMConfig, InsCodeModel
from ..mcps import MCPClientManager
from ..models.graph import GraphRequest
from ..parser.xml_parser import XmlToolResultParser
from ..config.agents import LLMType
from loguru import logger


class DesignerGraph(BaseGraph):
    """
    专业的UI/UX设计师和原型师AI助手，专注于根据用户需求生成设计图和原型。
    应用场景：界面设计、原型图制作、设计稿生成、用户体验优化、视觉设计、交互设计等所有与设计相关的工作。
    职责范围：
    - 界面设计：创建现代化、美观的用户界面设计图
    - 原型制作：制作交互原型图、线框图和设计稿
    - 视觉设计：色彩搭配、字体选择、布局设计和视觉层次的图像呈现
    - 用户体验：优化用户流程、提升交互体验的设计图
    - 设计系统：建立设计规范、组件库和设计标准的视觉图
    - 响应式设计：适配不同设备和屏幕尺寸的设计图方案
    - 移动端设计：iOS和Android应用界面设计图

    注意：本助手专注于生成设计图和原型图，不具备代码生成能力，仅专注于视觉设计和图像创作。
    """
    llm_vl: InsCodeModel = None
    llm_image: InsCodeModel = None

    def __init__(self,
                 llm_vl: InsCodeModel | LLMConfig = None,
                 llm_image: InsCodeModel | LLMConfig = None,
                 redis: Redis = None,
                 nodes: List[BaseNodes] = [],
                 graph_name="agent_design_",
                 agent_template: str = "agent_design",
                 mcp_client: MCPClientManager = None,
                 **kwargs):
        super().__init__(graph_name=graph_name,
                         agent_template=agent_template,
                         redis=redis,
                         mcp_client=mcp_client,
                         **kwargs)
        nodes = nodes if nodes else []

        self.llm_vl = InsCodeModel.build(llm_vl if llm_vl else LLMConfig.build_from_type(LLMType.LLM_VL))
        self.llm_image = InsCodeModel.build(llm_image if llm_image else LLMConfig.build_from_type(LLMType.LLM_IMAGE))

        self.nodes = nodes if nodes else [AskUserNodes(), CommandNodes(), CompletionNodes(), FileNodes(), WebNodes(),
                                          ImageNodes(llm_vl=self.llm_vl, llm_image=self.llm_image)]
        if mcp_client and mcp_client.is_available():
            self.nodes.append(MCPNodes(mcp_client))
        self.parser = XmlToolResultParser(tools=self.get_nodes_schema())

    async def get_system_template(self, request: GraphRequest, state=None, template_type: str = None):
        """
        Override base implementation to add coder-specific default values.
        """
        if "custom_instructions" not in request.kwargs:
            request.kwargs["custom_instructions"] = ""

        return await super().get_system_template(request=request, state=state, template_type=template_type)

    async def destroy(self):
        """
        Clean up designer-specific resources and call parent cleanup.
        """
        try:
            # Clean up vision LLM client
            if self.llm_vl:
                await self.llm_vl.close()
                self.llm_vl = None

            if self.llm_image:
                await self.llm_image.close()
                self.llm_image = None

            await super().destroy()
        except Exception as e:
            logger.error(f"Error during designer graph cleanup: {e}")
